#!/usr/bin/env python3
import asyncio
import logging
from datetime import datetime
from typing import Any, Dict
from mcp.server.fastmcp import FastMCP, Context
from mcp.server.sse import SseServerTransport
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.routing import Mount, Route
from starlette.responses import PlainTextResponse
import uvicorn

# 简单日志配置
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
logger = logging.getLogger("mcp")

# 创建FastMCP服务器 (24小时超时)
mcp = FastMCP("progress-test-server", server_options={"timeout": 86400})

# 全局进度追踪
current_progress = 0
start_time = None


@mcp.tool()
async def progress_test(iterations: int = 100000, sleep_time: float = 0.1, ctx: Context = None) -> Dict[str, Any]:
    """持续调用report_progress测试最长支持时间"""
    global current_progress, start_time
    start_time = datetime.now()
    current_progress = 0

    if not ctx:
        return {"error": "无上下文对象"}

    try:
        for i in range(iterations):
            current_progress = i + 1

            # 仅在关键节点记录日志，减少噪音
            if current_progress % 100 == 0:
                elapsed = (datetime.now() - start_time).total_seconds()
                logger.info(f"进度: {current_progress}/{iterations} - 已运行: {elapsed:.1f}秒")

            # 持续报告进度
            await ctx.report_progress(float(current_progress), float(iterations))
            await asyncio.sleep(sleep_time)

    except Exception as e:
        return {"error": str(e), "iterations_completed": current_progress}

    total_time = (datetime.now() - start_time).total_seconds()
    return {"completed": current_progress, "seconds": total_time}


# 创建服务器应用
def create_app():
    sse = SseServerTransport("/messages/")

    async def handle_sse(request: Request):
        try:
            async with sse.connect_sse(request.scope, request.receive, request._send) as (r, w):
                await mcp._mcp_server.run(r, w, mcp._mcp_server.create_initialization_options())
                return PlainTextResponse("Connection closed")
        except Exception as e:
            return PlainTextResponse(f"Error: {str(e)}", status_code=500)

    return Starlette(
        routes=[
            Route("/sse", endpoint=handle_sse),
            Mount("/messages/", app=sse.handle_post_message),
        ],
    )


if __name__ == "__main__":
    print("\n===== 进度报告测试服务器 =====")
    print("端点: http://localhost:8000/sse")
    print("客户端: npx -y @modelcontextprotocol/inspector http://localhost:8000/sse --timeout 86400000")
    uvicorn.run(create_app(), host="localhost", port=8000)